% Updates G
%
% Jon Steinsson and Emi Nakamura, August 2006
%****************************************************

function [G_new W_SPold supDifG G_temp GMove] ...
    = UpdateGCalvoPlus(Q,F,F2,G_old,...
    a,rp,rad,alphaCalvo,theta,Proba,ProbNgr,GMove_old)

agridsize = size(a,1);
rpgridsize = size(rp,1);
radgridsize = size(rad,1);
Ssize = agridsize*rpgridsize*radgridsize;

Q = reshape(Q,[Ssize 1]);
F = reshape(F,[Ssize 1]);
F2 = reshape(F2,[Ssize 1]);

% Get weights as a function of S_t/P_t-1
%********************************************************

QaInd = repmat((1:agridsize)',rpgridsize*radgridsize,1);

QradInd = repmat((1:radgridsize)',[1 agridsize rpgridsize]);
QradInd = permute(QradInd,[2 3 1]);
QradInd = reshape(QradInd,[Ssize 1]);

[dummy,QrpInd] = ismember(int32(F*10^6),int32(rp*10^6));
clear dummy
FInd = sub2ind([agridsize rpgridsize radgridsize],QaInd,QrpInd,QradInd);

[dummy,QrpInd] = ismember(int32(F2*10^6),int32(rp*10^6));
clear dummy
F2Ind = sub2ind([agridsize rpgridsize radgridsize],QaInd,QrpInd,QradInd);
clear QaInd QrpInd QradInd

OldInd = (1:Ssize)';

[FInd iFInd] = sort(FInd);
OldFInd = OldInd(iFInd);

[F2Ind iF2Ind] = sort(F2Ind);
OldF2Ind = OldInd(iF2Ind);

Q_old = Q;
% Apply F
Q = accumarray(FInd,Q_old(OldFInd));
Q = [Q; zeros(size(Q_old,1)-size(Q,1),1)];
Q2 = accumarray(F2Ind,Q_old(OldF2Ind));
Q2 = [Q2; zeros(size(Q_old,1)-size(Q2,1),1)];
Q = alphaCalvo*Q + (1-alphaCalvo)*Q2;

% Apply Proba
Q = ((reshape(Q,agridsize,rpgridsize*radgridsize)')*Proba)';

% Apply ProbNgr
Q = reshape(Q,[agridsize*rpgridsize radgridsize])*ProbNgr;
Q = reshape(Q,[Ssize 1]);

% Get the prices
%**********************************************************

GInd = CalculateGInd(a,rp,rad,G_old);

F = F(GInd);
F2 = F2(GInd);

%figure(22)
%F_plot = reshape(F,[agridsize rpgridsize radgridsize]);
%F_plot = squeeze(F_plot(10,:,:));
%surf(exp(rad),exp(rp),exp(F_plot))
%title('Policy Function');
%xlabel('exp(rad)')
%ylabel('exp(rp)')

% Create Price Index
%***************************************************************

F = exp(F).^(1-theta);
F2 = exp(F2).^(1-theta);

F = alphaCalvo*F + (1-alphaCalvo)*F2;

totWeight = sum(reshape(Q,[agridsize*rpgridsize radgridsize]))';
F = reshape(F.*Q,[agridsize*rpgridsize radgridsize]);

% Deal with zero mass
W_zero = (totWeight < 10^(-10));
totWeight = totWeight + W_zero;

sumF = sum(F)';
sumF_zero = (sumF < 10^(-10));
sumF = sumF + sumF_zero;

G_new = log((sumF./totWeight).^(1/(1-theta)));
G_new = (1-sumF_zero).*(1-W_zero).*G_new ...
    + W_zero.*zeros(radgridsize,1) ...
    + (1-W_zero).*sumF_zero.*zeros(radgridsize,1);
G_temp = G_new;
supDifG = max(abs(G_new));

rpinc = rp(2) - rp(1);
GstepH = zeros(radgridsize,1);
GstepH(2:end,1) = (abs(G_old(1:end-1,1)-G_old(2:end,1)) >  rpinc/2);
GstepH(1) = 1;

GstepL = zeros(radgridsize,1);
GstepL(1:end-1,1) = (abs(G_old(1:end-1,1)-G_old(2:end,1)) >  rpinc/2);
GstepL(end) = 1;

GmoveU = (G_new > rpinc/1.6).*GstepL;
GmoveD = (G_new < - rpinc/1.6).*GstepH;

GmoveUshift = [0; GmoveU(1:end-1,1)];
GmoveDshift = [GmoveD(2:end,1); 0];

% Create indicators for the two cases
MoveUp = (GmoveU & (1-GmoveDshift)); 
MoveDw = (GmoveD & (1-GmoveUshift));

% Get rid of cycling on many points
while sum(MoveUp) > radgridsize/5
    MoveUp = MoveUp.*(rand(radgridsize,1) > 0.25);
end
while sum(MoveDw) > radgridsize/5
    MoveDw = MoveDw.*(rand(radgridsize,1) > 0.25);
end

%[G_old G_new GmoveU GmoveD MoveUp MoveDw]

GMove = MoveUp*rpinc - MoveDw*rpinc;
% Get rid of cycling on few points
if max(abs(GMove+GMove_old)) < rpinc/2
    GMove = zeros(radgridsize,1);
end

G_new = G_old + GMove;

%supDifG = max(abs(G_new-G_old));

if max(abs(G_new-G_old)) < rpinc/2
    supDifG = 0;
end


% rpinc = rp(2) - rp(1);
% if StageForG == 1
%     %G_new = G_old + 1.01*rpinc*(G_new.*rand(radgridsize,1) > rpinc) ...
%     %    - 1.01*rpinc*(G_new.*rand(radgridsize,1) < -rpinc);
%     G_new = G_old + 1.01*min(abs([rpinc*ones(1,radgridsize); G_new']))'.*(G_new.*rand(radgridsize,1) > rpinc/1.4) ...
%         - 1.01*min(abs([rpinc*ones(1,radgridsize); G_new']))'.*(G_new.*rand(radgridsize,1) < -rpinc/1.4);
% else
%     %[(G_new > rpinc) (G_new < -rpinc) (abs(G_new) == max(abs(G_new))*ones(size(G_new,1),1)) G_new]
%     G_new = G_old + 1.02*min(rpinc,max(abs(G_new)))*(G_new > 0).*(abs(G_new) == max(abs(G_new))*ones(size(G_new,1),1)) ...
%         - 1.02*min(rpinc,max(abs(G_new)))*(G_new < 0).*(abs(G_new) == max(abs(G_new))*ones(size(G_new,1),1));
% end
% 
% if StageForG == 1
%     supDifG = max(abs(G_new-G_old));
%     if supDifG < rpinc
%         %supDifG = max(abs(G_new-G_old));
%         supDifG = 2*rpinc;
%         StageForG = 2;
%     end
% else
%     supDifG = max(abs(G_new-G_old));
% end

% Put G_new on the grid 
G_new_mod = mod(G_new,rpinc);
G_new_mod_up = (G_new_mod > rpinc/2);
G_new = G_new + G_new_mod_up.*(rpinc-G_new_mod) ...
    - (1-G_new_mod_up).*G_new_mod;

% Create distribution of S_t/P_t-1
%***************************************************

W_SPold = sum(reshape(Q,[agridsize*rpgridsize radgridsize]))';
W_SPold = (W_SPold >= 10^(-10)).*W_SPold;

